import argparse
import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Union

import numpy as np
import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader
from torchvision import transforms

from yms_zsl.dataset_utils.CustomImageDataset import ImageClassificationDataset
from yms_zsl.models.HSAZLM import FCNN, LSELoss
from yms_zsl.tools.tool import create_csv, append_metrics_to_csv, get_device
from yms_zsl.tools.tools import setup_logger
from yms_zsl.tools.train_eval_utils import train_fcnn_one_epoch

logger = logging.getLogger('yms')

def train_fcnn(configs):
    device = configs.device
    transform = configs.transform
    save_dir = configs.save_dir
    fc_results_file = os.path.join(save_dir, 'fc_metrics.csv')
    fc_metrics = ['epoch', 'train_loss', 'val_loss', 'lr']
    create_csv(fc_metrics, fc_results_file)

    train_dataset = ImageClassificationDataset(configs.data_dir, transform=transform, train_class=configs.train_class,
                                               factor_index_map_path=configs.factor_index_map_path,
                                               ignore_factors=['Operating Condition'])
    val_dataset = ImageClassificationDataset(configs.data_dir, transform=transform, mode='val',
                                             train_class=configs.train_class, factor_index_map_path=configs.factor_index_map_path,
                                             ignore_factors=['Operating Condition']
                                             )
    train_loader = DataLoader(train_dataset, batch_size=configs.batch_size, shuffle=True,
                              num_workers=configs.num_workers)
    val_loader = DataLoader(val_dataset, batch_size=configs.batch_size, shuffle=False, num_workers=configs.num_workers)
    classes = train_dataset.classes

    semantic_attr_list = []
    for cls_name in classes:
        npy_path = os.path.join(save_dir, 'attributes', 'semantic_attribute', f'{cls_name}.npy')
        cls_attr = np.load(npy_path, allow_pickle=True)  # 读取单个类别特征
        semantic_attr_list.append(cls_attr)

    hsa_matrix = torch.tensor(np.array(semantic_attr_list, dtype=np.float32), dtype=torch.float32)

    model = FCNN(configs.fcnn_channels, os.path.join(save_dir, 'checkpoints', 'cnn.pth')).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=configs.lr)
    lr_scheduler = ReduceLROnPlateau(optimizer, factor=0.1, min_lr=1e-10, patience=7, mode='min')
    criterion = LSELoss(hsa_matrix).to(device)
    best_loss = 1e8
    num_epochs = configs.epochs
    # 添加早停相关参数
    patience = 10  # 从外部参数获取耐心值
    early_stop_counter = 0  # 早停计数器
    best_epoch = 0  # 最佳模型的epoch

    for epoch in range(num_epochs):
        training_lr = lr_scheduler.get_last_lr()[0]
        train_loss, val_loss = train_fcnn_one_epoch(
            model=model, optimizer=optimizer, train_loader=train_loader,
            val_loader=val_loader, device=device, criterion=criterion, epoch=epoch
        )
        lr_scheduler.step(val_loss)
        metric = {'epoch': epoch, 'train_loss': train_loss, 'val_loss': val_loss, 'lr': training_lr}
        append_metrics_to_csv(metric, fc_results_file)
        logger.info(f'val epoch {epoch + 1}, val metric: {metric}')
        current_score = val_loss
        if current_score < best_loss:
            best_loss = current_score
            best_epoch = epoch
            early_stop_counter = 0  # 重置早停计数器
            torch.save(model.state_dict(), os.path.join(save_dir, 'checkpoints', 'FCnn.pth'))
            logger.info(f'Best model saved at epoch {epoch + 1} with val_loss: {best_loss:.6f}')
        elif patience > 0 and training_lr < (configs.lr *0.001):  # 仅当patience>0时执行早停计数
            early_stop_counter += 1
            logger.info(f'Early stopping counter: {early_stop_counter}/{patience}')

            # 如果早停计数器达到耐心值，停止训练
            if early_stop_counter >= patience:
                logger.info(f'Early stopping triggered after {epoch + 1} epochs')
                break

    logger.info(f'Best model was saved at epoch {best_epoch + 1} with val_loss: {best_loss:.4f}')


def parse_args(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--epochs', type=int, default=500)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--fcnn_channels', type=int, default=518)
    parser.add_argument('--data_dir', default=r'/data/coding/data/D0')
    parser.add_argument('--output_dir', default='/data/coding/results/train_D0')
    parser.add_argument('--feature_path', default=r'/data/coding/results/train_D0/models/feature_extractor.pt')
    parser.add_argument('--hsa_path', default=r'/data/coding/results/train_D0/HSA.mat')
    return parser.parse_args(args if args else [])

@dataclass
class TrainFCnnConfig:
    # 设备与数据变换（添加明确类型注解）
    device: torch.device = field(
        default_factory=get_device,  # 用default_factory延迟初始化，避免模块导入时执行
        metadata={"desc": "训练使用的设备（CPU/GPU）"},
        repr=False

    )
    transform: transforms.Compose = field(
        default_factory=lambda: transforms.Compose([  # 用default_factory创建Compose实例
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]),
        metadata={"desc": "图像预处理变换管道"},
        repr=False
    )
    # -------------------------- 路径参数--------------------------
    data_dir: Union[str, Path] = r'D:\Code\2-ZSL\0-data\CWRU\dataset'  # 扩展为Union类型
    save_dir: Union[str, Path] = r'D:\Code\2-ZSL\1-output\特征解耦结果\exp-3'
    train_class: Union[str, Path] = r'D:\Code\2-ZSL\0-data\CWRU\dataset\seen_classes.txt'
    factor_index_map_path: Union[str, Path] = None

    # -------------------------- 训练超参数--------------------------
    epochs: int = 1000
    batch_size: int = 40
    lr: float = 1e-3
    weight_decay: float = 1e-5
    patience: int = 10
    num_workers: int = 0
    fcnn_channels: int = 518

if __name__ == '__main__':
    opts = TrainFCnnConfig(
        data_dir=r'D:\Code\2-ZSL\0-data\data\data',
        save_dir=r'D:\Code\2-ZSL\1-output\论文实验结果\对比方法\CWRU\HSAZLM\H06\exp-1',
        train_class=r'D:\Code\2-ZSL\0-data\data\data/seen_classes.txt',
        batch_size=256,
        factor_index_map_path=r'D:\Code\2-ZSL\0-data\data\data\factor_index_map.txt'
    )
    logger = setup_logger(opts.save_dir)
    logger.info("=" * 60)
    logger.info("【单独运行】FCnn训练流程")
    logger.info(f"配置信息:{vars(opts)}")
    logger.info("=" * 60)
    train_fcnn(opts)
